﻿module EmitEnvironment

open Ast
open EmitType
open StateMonad
open SymbolTableTypes
open System.Collections.Generic

type Environment = 
  { Emit : List<Emit>
//    tob : int
    mutable ActualTable   : BlockTable
    L_counter             : int
    T_Symbols             : List<T_symbol>
    mutable While_counter : int
    mutable Actual_while  : int list }


type State<'a> = Environment -> ('a * Environment)

let private get_table =
  fun s ->
    (s.ActualTable, s)

let rec mapM pred xs =
  state{
    match xs with
    | [] -> 
      return []
    | x :: xs ->
      let! x = pred x
      let! xs = mapM pred xs
      return (x :: xs)
  }

let emit (i : Emit) =
  fun s ->
    s.Emit.Add i
    ((), s)

let new_label : State<EmitLabel> =
  fun s ->
    let s' = { s with L_counter  = s.L_counter + 1 }
    (ref (sprintf "Label%d" s'.L_counter), s')

let new_while_label : State<EmitLabel * EmitLabel> =
  fun s ->
    do s.While_counter <- s.While_counter + 1
    do s.Actual_while <- s.While_counter :: s.Actual_while
    ((ref(sprintf "StartWhile%d" s.While_counter),
      ref(sprintf "EndWhile%d" s.While_counter)), s)

let get_last_start_while_label : State<EmitLabel> =
  fun s ->
    (ref(sprintf "StartWhile%d" s.Actual_while.Head), s)

let get_last_end_while_label : State<EmitLabel> =
  fun s ->
    (ref(sprintf "EndWhile%d" s.Actual_while.Head), s)

let exit_while =
  fun s ->
    let s' = { s with Actual_while  = s.Actual_while.Tail }
    (), s'

let add_label =
  fun s ->
    s.Emit.Add(Emit.Label (ref(sprintf "label%d" s.L_counter)))
    ((), s)

let private new_t = 
  fun s ->
    let t = { IsUsed = true; Id = s.T_Symbols.Count; IsArg = false } 
    do s.T_Symbols.Add(t)
    (t, s)

let get_t = 
  fun (s : Environment) ->
    match s.T_Symbols |> Seq.tryFind (fun t_s -> not t_s.IsUsed) with
    | None -> 
      new_t s
    | Some var ->
      do var.IsUsed <- true
      (var, s)

let free_t (t : T_symbol) =
  fun (s : Environment) ->
    match s.T_Symbols |> Seq.tryFind (fun t_s -> t_s.Name = t.Name) with
    | None -> (), s
    | Some var ->
      var.IsUsed <- false
      (), s

let enter_block block =
  fun (s : Environment) ->
    if s.ActualTable.Block = block then
      ((), s)
    else
      do s.ActualTable <- s.ActualTable.Children.Head
      ((), s)

let exit_block = 
  fun (s : Environment) ->
    match s.ActualTable.Parent with
    | BodyOrBlock.Body b ->
      ((), s)
    | BodyOrBlock.Block b ->
      do s.ActualTable <- b
      do b.Children.Enqueue(b.Children.Pop())
      ((), s)

let rec try_find_var name = 
  state{
    let! first_env = get_table
    let rec exists (t : GenericTable) = 
      state{
        match t with
        | GenericTable.BodyTable b ->
          if b.Args.ContainsKey(name) then
            return Some b.Args.[name]
          else
            return None
        | GenericTable.BlockTable b ->
          if b.Symbols.ContainsKey(name) then
            return Some b.Symbols.[name]
          else
            return! exists (match b.Parent with | BodyOrBlock.Block b -> GenericTable.BlockTable b | BodyOrBlock.Body b -> GenericTable.BodyTable b )
      }
    return! exists (GenericTable.BlockTable first_env)
  }

let get_class name =
  state{
    let! t = get_table
    let class_table = get_class_table (GenericTable.BlockTable t)
    let global_table = get_global_table (GenericTable.ClassTable class_table)
    return global_table.[name]
  }

let get_container_class = 
  state{
    let! state = getState
    let block_table = GenericTable.BlockTable state.ActualTable
    return (get_class_table block_table)
  }

let make_object_layout class_table =
  let tables =  get_inherited_class_table (GenericTable.ClassTable class_table) |> List.rev
  let fields_list = ResizeArray()
  do for f_or_m in tables do
      f_or_m.Fields 
                  |> Seq.iter(fun key_value -> 
                                  if fields_list |> Seq.exists((=) key_value.Key) |> not then
                                    fields_list.Add (key_value.Key))
  do class_table.ObjectLayout.AddRange(fields_list)

let make_dispatch_vector class_table =
  let tables =  get_inherited_class_table (GenericTable.ClassTable class_table) |> List.rev
  let fields_list = ResizeArray()
  do for f_or_m in tables do
      f_or_m.Methods 
                  |> Seq.iter(fun key_value -> 
                                  if fields_list |> Seq.exists(fun (name, _) -> name = key_value.Key) |> not then
                                    fields_list.Add (key_value.Key, ref f_or_m.Name)
                                  else 
                                    let _, owner_class = fields_list |> Seq.find(fun (name, _) -> name = key_value.Key)
                                    do owner_class := f_or_m.Name)
  do class_table.DispatchVector.AddRange(fields_list)


let given_class_get_field_index name (class_table : ClassTable) =
  class_table.ObjectLayout |> Seq.findIndex((=) name)


let get_field_index name : State<int, Environment> =
  state{
    let! class_table = get_container_class
    return given_class_get_field_index name class_table
  }

let get_method_index name _class : State<int, Environment> =
  state{    
    let! state = getState
    let block_table = GenericTable.BlockTable state.ActualTable

    let class_table = 
      if _class =  None then        
        get_class_table block_table
      else 
        let (Identifier _class) = _class.Value
        (get_global_table block_table).[_class]


    return class_table.DispatchVector |> Seq.findIndex(fun (method_name, _) -> method_name = name)
  }

let get_t_qualified name =
  state{
    let! maybe_var = try_find_var name
    let! s = getState
    match maybe_var with
    | Some var ->
      let body_env = get_body_table (GenericTable.BlockTable s.ActualTable)
      let alias = body_env.Aliases |> Seq.find (fun (a : BodySymbol) -> a.Alias = var.Alias)
      let ordinal = alias.Alias.Split([|'$'|]).[1] |> int
      return s.T_Symbols.[ordinal]
    | None ->
      let! t = get_t
      let! layout_index = get_field_index name
      do! emit (Emit.Expr (EmitExpr.GetField(s.T_Symbols.[0], layout_index), t))
      return t
  }



let set_t_qualified name t_val =
  state{
    let! maybe_var = try_find_var name
    let! s = getState
    match maybe_var with
    | Some var ->
      let body_env = get_body_table (GenericTable.BlockTable s.ActualTable)
      let alias = body_env.Aliases |> Seq.find (fun (a : BodySymbol) -> a.Alias = var.Alias)
      let ordinal = alias.Alias.Split([|'$'|]).[1] |> int
      do! emit (Emit.Stmt (EmitStmt.Assign (s.T_Symbols.[ordinal], t_val)))
    | None ->
      let! layout_index = get_field_index name      
      do! emit (Emit.Stmt (EmitStmt.SetField(s.T_Symbols.[0], layout_index, t_val)))
  }

let comment str = 
  state{
    do! emit (Emit.Comment ("#" + str))
  }

//let given_expr_get_field_index (expr : Expr) name =
//  state{
//      match expr with
//      | Expr.This -> 
//        return! get_field_index name
//      | Expr.New class_name -> 
//        let! actual_class = get_container_class
//        let class_table = (get_global_table (GenericTable.ClassTable actual_class)).[class_name]
//        return! given_class_get_field_index name class_table
//      | Expr.Call call -> 
//        let! ret = check_call call
//        match ret with
//        | ReturnType.Type ty -> 
//          return ty
//        | ReturnType.Void ->
//          return! fail "expression call must be non void"      
//  }

//      


//      | Expr.Location loc -> 
//        return! check_loc loc